import rospy
import sys
import tf2_ros as tf2
import tf_conversions as tfc
from geometry_msgs.msg import PoseStamped, TransformStamped
from tf2_msgs.msg import TFMessage
import math
import numpy
from nav_msgs.msg import Odometry

uav_num = 6
cur_state = [PoseStamped() for i in range(uav_num)]
def current_state_callback(data,i):
    cur_state[i] = data
    # print('upward vel %.2f', cur_state.pose.pose.orientation.x)

def qRot(p, q):
    qInvert = tfc.transformations.quaternion_inverse
    qMultiply = tfc.transformations.quaternion_multiply
    qInv = qInvert(q)
    pNew = qMultiply(qMultiply(q, p), qInv)
    return pNew

def qRotInv(p, q):
    qInvert = tfc.transformations.quaternion_inverse
    qMultiply = tfc.transformations.quaternion_multiply
    qInv = qInvert(q)
    pNew = qMultiply(qMultiply(qInv, p), q)
    return pNew


def __genTf(frameS, frameT, posTs, q):
    ts = TransformStamped()
    ts.header.stamp = rospy.Time.now()
    ts.header.frame_id = frameS
    ts.child_frame_id = frameT
    
    ts.transform.translation.x = posTs[0]
    ts.transform.translation.y = posTs[1]
    ts.transform.translation.z = posTs[2]

    ts.transform.rotation.x = q[0]
    ts.transform.rotation.y = q[1]
    ts.transform.rotation.z = q[2]
    ts.transform.rotation.w = q[3]
    return ts

def __main():
    rospy.init_node("tf_broadcaster")
    # rate = rospy.Rate(1)
    # tbf = tf2.Buffer()
    # tl = tf2.TransformListener(tbf)
    # bias_z = [0.,0.38,0.5]
    # a = 0
    for i in range(3):
        rospy.Subscriber("/UAV"+str(i)+"/mavros/local_position/pose", PoseStamped, current_state_callback, i, queue_size = 2)
    for i in range(3,uav_num):
        rospy.Subscriber("/UAV"+str(i)+"/global_position/uav_data", PoseStamped, current_state_callback, i,queue_size = 2)
    roundCount = 0
    rateFast = rospy.Rate(15)
    qInvert = tfc.transformations.quaternion_inverse
    qMultiply = tfc.transformations.quaternion_multiply
    while not rospy.is_shutdown():
        print("begin: "+ str(roundCount))
        roundCount += 1

        bc = tf2.TransformBroadcaster()
        # p02t = numpy.array([2 * math.cos(a), -1 + 2 * math.sin(a), 2 + math.sin(5 * a), 0])
        # q02t = numpy.array(tfc.transformations.quaternion_from_euler(0, 0, math.pi / 2 + a))  # math.pi / 2 +
        # ts = __genTf("state_0", "target_state", p02t, q02t)
        # bc.sendTransform(ts)
        # if a < math.pi * 2:
        #     a += 0.01
        # else:
        #     a = 0

        for i in range(uav_num):
            if cur_state[i].pose.orientation.w != 0.0:
            # # 02c
            #     if i < 3:
            #         z_bias = bias_z[i]
            #     else:
            #         z_bias = 0
                p02c = numpy.array([cur_state[i].pose.position.x, cur_state[i].pose.position.y, cur_state[i].pose.position.z, 0])
                q02c = (cur_state[i].pose.orientation.x, cur_state[i].pose.orientation.y, cur_state[i].pose.orientation.z, cur_state[i].pose.orientation.w)
            # print("currently:\t x %.2f\t y %.2f\t z %.2f" % (
            #     cur_state.pose.pose.position.x, cur_state.pose.pose.position.y, cur_state.pose.pose.position.z))
            #     if i == 5:
            #         print(p02c)
                ts = __genTf("state_0", "UAV_"+str(i), p02c, q02c)
                # print i
                bc.sendTransform(ts)

        rateFast.sleep()


if __name__ == "__main__":
    __main() 
